{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from azure.ai.projects import AIProjectClient\n",
    "from azure.ai.projects.models import CodeInterpreterTool\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from typing import Any\n",
    "from pathlib import Path\n",
    "from datetime import datetime\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project_client = AIProjectClient.from_connection_string(\n",
    "    credential=DefaultAzureCredential(), conn_str=os.environ[\"PROJECT_CONNECTION_STRING\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, HTML, Image\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "async def run_agent_with_visualization():\n",
    "    html_output = \"<h2>Azure AI Agent Execution</h2>\"\n",
    "\n",
    "    with project_client:\n",
    "        # Create an instance of the CodeInterpreterTool\n",
    "        code_interpreter = CodeInterpreterTool()\n",
    "\n",
    "        # The CodeInterpreterTool needs to be included in creation of the agent\n",
    "        agent = project_client.agents.create_agent(\n",
    "            model=\"gpt-4o-mini\",\n",
    "            name=\"my-agent\",\n",
    "            instructions=\"You are helpful agent\",\n",
    "            tools=code_interpreter.definitions,\n",
    "            tool_resources=code_interpreter.resources,\n",
    "        )\n",
    "        html_output += f\"<div><strong>Created agent</strong> with ID: {agent.id}</div>\"\n",
    "\n",
    "        # Create a thread\n",
    "        thread = project_client.agents.create_thread()\n",
    "        html_output += f\"<div><strong>Created thread</strong> with ID: {thread.id}</div>\"\n",
    "\n",
    "        # User query - display nicely\n",
    "        user_query = \"Could you please create a bar chart for the operating profit using the following data and provide the file to me? Bali: 100 Travelers, Paris: 356 Travelers, London: 900 Travelers, Tokyo: 850 Travellers\"\n",
    "        html_output += \"<div style='margin:15px 0; padding:10px; background-color:#f5f5f5; border-left:4px solid #007bff; border-radius:4px;'>\"\n",
    "        html_output += \"<strong>User:</strong><br>\"\n",
    "        html_output += f\"<div style='margin-left:15px'>{user_query}</div>\"\n",
    "        html_output += \"</div>\"\n",
    "\n",
    "        # Create a message\n",
    "        message = project_client.agents.create_message(\n",
    "            thread_id=thread.id,\n",
    "            role=\"user\",\n",
    "            content=user_query,\n",
    "        )\n",
    "\n",
    "        # Run the agent - show a \"processing\" message\n",
    "        display(HTML(\n",
    "            html_output + \"<div style='color:#007bff'><i>Processing request...</i></div>\"))\n",
    "\n",
    "        # Execute the run\n",
    "        run = project_client.agents.create_and_process_run(\n",
    "            thread_id=thread.id, assistant_id=agent.id)\n",
    "\n",
    "        # Update status\n",
    "        status_color = 'green' if run.status == 'completed' else 'red'\n",
    "        html_output += f\"<div><strong>Run finished</strong> with status: <span style='color:{status_color}'>{run.status}</span></div>\"\n",
    "\n",
    "        if run.status == \"failed\":\n",
    "            html_output += f\"<div style='color:red'><strong>Run failed:</strong> {run.last_error}</div>\"\n",
    "\n",
    "        # Get messages from the thread\n",
    "        messages = project_client.agents.list_messages(thread_id=thread.id)\n",
    "\n",
    "        # Format assistant response\n",
    "        html_output += \"<div style='margin:15px 0; padding:10px; background-color:#f0f7ff; border-left:4px solid #28a745; border-radius:4px;'>\"\n",
    "        html_output += \"<strong>Assistant:</strong><br>\"\n",
    "\n",
    "        # Handle messages based on the actual structure\n",
    "        # First, try to get the assistant's text responses\n",
    "        try:\n",
    "            # First approach - if messages is a list of objects with role attribute\n",
    "            assistant_msgs = [msg for msg in messages if hasattr(\n",
    "                msg, 'role') and msg.role == \"assistant\"]\n",
    "\n",
    "            if assistant_msgs:\n",
    "                last_msg = assistant_msgs[-1]\n",
    "                if hasattr(last_msg, 'content'):\n",
    "                    if isinstance(last_msg.content, list):\n",
    "                        for content_item in last_msg.content:\n",
    "                            if hasattr(content_item, 'type') and content_item.type == \"text\":\n",
    "                                html_output += f\"<div style='margin-left:15px; white-space:pre-wrap'>{content_item.text.value}</div>\"\n",
    "                    elif isinstance(last_msg.content, str):\n",
    "                        html_output += f\"<div style='margin-left:15px; white-space:pre-wrap'>{last_msg.content}</div>\"\n",
    "\n",
    "            # If no messages were found with the above approach, try a different structure\n",
    "            if not assistant_msgs:\n",
    "                # If messages is a class with attributes\n",
    "                if hasattr(messages, 'data'):\n",
    "                    for msg in messages.data:\n",
    "                        if hasattr(msg, 'role') and msg.role == \"assistant\":\n",
    "                            if hasattr(msg, 'content'):\n",
    "                                html_output += f\"<div style='margin-left:15px; white-space:pre-wrap'>{msg.content}</div>\"\n",
    "\n",
    "        except Exception as e:\n",
    "            html_output += f\"<div style='color:red'><strong>Error processing messages:</strong> {str(e)}</div>\"\n",
    "\n",
    "        html_output += \"</div>\"\n",
    "\n",
    "        # Handle image contents based on the actual structure\n",
    "        saved_images = []\n",
    "        try:\n",
    "            # Try to access image_contents as an attribute\n",
    "            if hasattr(messages, 'image_contents'):\n",
    "                for image_content in messages.image_contents:\n",
    "                    file_id = image_content.image_file.file_id\n",
    "                    file_name = f\"{file_id}_image_file.png\"\n",
    "                    project_client.agents.save_file(\n",
    "                        file_id=file_id, file_name=file_name)\n",
    "                    saved_images.append(file_name)\n",
    "                    html_output += f\"<div style='margin-top:10px'><strong>Generated Image:</strong> {file_name}</div>\"\n",
    "        except Exception as e:\n",
    "            html_output += f\"<div style='color:orange'><i>Note: No images found or error processing images</i></div>\"\n",
    "\n",
    "        # Handle file path annotations based on the actual structure\n",
    "        try:\n",
    "            # Try to access file_path_annotations as an attribute\n",
    "            if hasattr(messages, 'file_path_annotations'):\n",
    "                for file_path_annotation in messages.file_path_annotations:\n",
    "                    file_name = Path(file_path_annotation.text).name\n",
    "                    project_client.agents.save_file(\n",
    "                        file_id=file_path_annotation.file_path.file_id, file_name=file_name)\n",
    "                    html_output += \"<div style='margin:10px 0; padding:8px; background-color:#f8f9fa; border:1px solid #ddd; border-radius:4px;'>\"\n",
    "                    html_output += f\"<strong>Generated File:</strong> {file_name}<br>\"\n",
    "                    html_output += f\"<strong>Type:</strong> {file_path_annotation.type}<br>\"\n",
    "                    html_output += \"</div>\"\n",
    "        except Exception as e:\n",
    "            html_output += f\"<div style='color:orange'><i>Note: No file annotations found or error processing files</i></div>\"\n",
    "\n",
    "        # Delete the agent once done\n",
    "        project_client.agents.delete_agent(agent.id)\n",
    "        html_output += \"<div style='margin-top:10px'><i>Agent deleted after completion</i></div>\"\n",
    "\n",
    "        # Final display of all content\n",
    "        display(HTML(html_output))\n",
    "\n",
    "        # Display any saved images\n",
    "        for img_file in saved_images:\n",
    "            display(Image(img_file))\n",
    "\n",
    "# Execute the function\n",
    "await run_agent_with_visualization()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
